#gui
import tkinter as tk
import matplotlib.pyplot as plt

class RealTimePlotter:
    '''实时绘图'''

    def __init__(self):
        self.tk_root = tk.Tk()
        self.tk_root.withdraw()        
        self.fig, (self.ax1, self.ax2) = plt.subplots(2, 1, figsize=(12, 8))
        self.ax1.set_ylabel('电压 (V)')
        self.ax2.set_ylabel('预紧力 (N)')
        self.ax2.set_xlabel('训练轮次')
        self.ax1.set_title('实测电压 vs 理论电压')
        self.ax2.set_title('预测预紧力')
        self.V_real_line, = self.ax1.plot([], [], 'bo-', label='实测电压')
        self.V_pred_line, = self.ax1.plot([], [], 'r--', label='理论电压')
        self.F_pred_line, = self.ax2.plot([], [], 'g-o', label='预测预紧力')
        self.ax1.legend()
        self.ax2.legend()
        self.epoch_data = []
        self.V_real_data = []
        self.V_pred_data = []
        self.F_pred_data = []
        plt.ion() 
        plt.show(block=False)
        
    def update(self, epoch, V_real, V_pred, F_pred):
        self.epoch_data.append(epoch)
        self.V_real_data.append(V_real)
        self.V_pred_data.append(V_pred)
        self.F_pred_data.append(F_pred)
        if len(self.epoch_data) > 100:
            for lst in [self.epoch_data, self.V_real_data, self.V_pred_data, self.F_pred_data]:
                lst.pop(0)        
        # 实时更新
        self.V_real_line.set_data(self.epoch_data, self.V_real_data)
        self.V_pred_line.set_data(self.epoch_data, self.V_pred_data)
        self.F_pred_line.set_data(self.epoch_data, self.F_pred_data)
        
        self.ax1.relim()
        self.ax1.autoscale_view()
        self.ax2.relim()
        self.ax2.autoscale_view()
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        plt.pause(0.001)  
        
    def close(self):
        plt.close(self.fig)
        self.tk_root.destroy()